package main

import (
	"errors"
	"fmt"
	"github.com/Shopify/sarama"
	cluster "github.com/bsm/sarama-cluster"
	"os"
	"os/signal"
)

func validate(group, topic string, brokers []string) (err error) {
	var (
		cli              *cluster.Client
		c                *cluster.Config
		broker           *sarama.Broker
		gresp            *sarama.DescribeGroupsResponse
		memberAssignment *sarama.ConsumerGroupMemberAssignment
		consumerNum      int
		partitions       []int32
	)
	c = cluster.NewConfig()
	c.Version = sarama.V0_10_0_1
	if cli, err = cluster.NewClient(brokers, c); err != nil {
		panic(err)
	}
	defer cli.Close()
	if partitions, err = cli.Partitions(topic); err != nil {
		return
	}
	if len(partitions) <= 0 {
		err = errors.New("partitions len <= 0")
		return
	}
	if err = cli.RefreshCoordinator(group); err != nil {
		return
	}
	if broker, err = cli.Coordinator(group); err != nil {
		return
	}
	defer broker.Close()

	if gresp, err = broker.DescribeGroups(&sarama.DescribeGroupsRequest{
		Groups: []string{group},
	}); err != nil {
		panic(err)
	}
	if len(gresp.Groups) != 1 {
		err = errors.New("gresp.Groups len != 1")
		return
	}
	for _, member := range gresp.Groups[0].Members {
		if memberAssignment, err = member.GetMemberAssignment(); err != nil {
			return
		}
		for mtopic := range memberAssignment.Topics {
			if mtopic == topic {
				consumerNum++
				break
			}
		}
	}

	// 这里不被允许使用已经分配给其他消费者的 topic
	if consumerNum >= len(partitions) {
		err = errors.New("useless consumer")
		return
	}
	return
}

func main() {
	// init (custom) config, set mode to ConsumerModePartitions
	config := cluster.NewConfig()
	config.Group.Mode = cluster.ConsumerModePartitions

	// init consumer
	brokers := []string{"127.0.0.1:9092"}
	topics := []string{"my_topic", "other_topic", "LiveDanmuSend-T"}
	group := "LiveDanmuSend-LiveLive-S"
	consumer, err := cluster.NewConsumer(brokers, group, topics, config)
	if err != nil {
		panic(err)
	}

	defer consumer.Close()

	if err := validate(group, topics[2], brokers); err != nil {
		panic(err)
	}
	// trap SIGINT to trigger a shutdown.
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, os.Interrupt)

	// consume partitions
	for {
		select {
		case part, ok := <-consumer.Partitions():
			if !ok {
				return
			}

			// start a separate goroutine to consume messages
			go func(pc cluster.PartitionConsumer) {
				for msg := range pc.Messages() {
					fmt.Fprintf(os.Stdout, "%s/%d/%d\t%s\t%s\n", msg.Topic, msg.Partition, msg.Offset, msg.Key, msg.Value)
					consumer.MarkOffset(msg, "") // mark message as processed
				}
			}(part)
		case <-signals:
			return
		}
	}

}
